from typing import Dict, Any
from .base_agent import BaseAgent, AgentResult
from .trending_tool import get_weibo_trending, get_toutiao_trending, get_douyin_trending
from datetime import datetime

class DataCollectionAgent(BaseAgent):
    """重构的数据收集智能体"""
    
    def __init__(self, config: Dict[str, Any] = None):
        super().__init__()
        self.config = config or {}
        self.name = "DataCollectionAgent"
        self.version = "1.0.0"
    
    def get_name(self) -> str:
        return self.name
    
    def get_version(self) -> str:
        return self.version
    
    def execute(self, input_data: Dict[str, Any]) -> AgentResult:
        self.log_execution_start()
        start_time = datetime.now()
        
        try:
            # 并行获取各平台数据
            weibo_data = get_weibo_trending()
            toutiao_data = get_toutiao_trending()
            douyin_data = get_douyin_trending()
            
            classified_data = {
                "weibo": {
                    "platform": "微博",
                    "data": weibo_data,
                    "count": len(weibo_data) if weibo_data else 0
                },
                "toutiao": {
                    "platform": "头条",
                    "data": toutiao_data,
                    "count": len(toutiao_data) if toutiao_data else 0
                },
                "douyin": {
                    "platform": "抖音",
                    "data": douyin_data,
                    "count": len(douyin_data) if douyin_data else 0
                }
            }
            
            total_count = sum([
                classified_data["weibo"]["count"],
                classified_data["toutiao"]["count"], 
                classified_data["douyin"]["count"]
            ])
            
            execution_time = (datetime.now() - start_time).total_seconds()
            self.log_execution_end(True, execution_time)
            
            return AgentResult(
                success=True,
                data=classified_data,
                metadata={
                    "total_count": total_count,
                    "platforms": ["weibo", "toutiao", "douyin"],
                    "timestamp": datetime.now().isoformat(),
                    "execution_time": execution_time
                }
            )
            
        except Exception as e:
            execution_time = (datetime.now() - start_time).total_seconds()
            self.log_execution_end(False, execution_time)
            return AgentResult(
                success=False,
                data={},
                error=f"数据收集失败: {str(e)}"
            )
